import uvicorn

from a2a.server.apps import A2AStarletteApplication
from a2a.server.request_handlers import DefaultRequestHandler
from a2a.server.tasks import InMemoryTaskStore
from a2a.types import (
    AgentCapabilities,
    AgentCard,
    AgentSkill,
)
from agent_executor import (
    WritingAgentExecutor,
)


def run_a2a_server():
    # 定义技能
    skill = AgentSkill(
        id='writing_skill',
        name='生成写作文章标题',
        description='生成写作文章标题',
        tags=['写作'],
        examples=['打破壁垒', '赋能协作', '保护隐私'],
    )

    extended_skill = AgentSkill(
        id='super_writing_skill',
        name='生成写作文章中心思想',
        description='只为认证用户提供生成写作文章中心思想的服务',
        tags=['写作'],
        examples=['仰望星空，脚踏实地', '打破壁垒、赋能协作、保护隐私'],
    )

    # 定义智能体卡片
    public_agent_card = AgentCard(
        name='智能写作 Agent - 普通用户版',
        description='智能写作 Agent - 普通用户版',
        url='http://localhost:9999/',
        version='1.0.0',
        defaultInputModes=['text'],
        defaultOutputModes=['text'],
        capabilities=AgentCapabilities(streaming=True),
        skills=[skill],
        supportsAuthenticatedExtendedCard=True,
    )

    specific_extended_agent_card = public_agent_card.model_copy(
        update={
            'name': '智能写作 Agent - 会员版',
            'description': '智能写作 Agent - 会员版',
            'version': '1.0.1',
            'skills': [
                skill,
                extended_skill,
            ],
        }
    )

    request_handler = DefaultRequestHandler(
        agent_executor=WritingAgentExecutor(),
        task_store=InMemoryTaskStore(),
    )

    server = A2AStarletteApplication(
        agent_card=public_agent_card,
        http_handler=request_handler,
        extended_agent_card=specific_extended_agent_card,
    )

    uvicorn.run(server.build(), host='0.0.0.0', port=9999)


if __name__ == '__main__':
    run_a2a_server()
